import torch
import torch.nn as nn
from einops import rearrange


class LSTMEncoder(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        
        self.net = nn.LSTM(input_size, hidden_size, batch_first=True)
        
        self.hour_embedding = nn.Embedding(30, hidden_size)
        self.day_embedding = nn.Embedding(40, hidden_size)
        self.month_embedding = nn.Embedding(20, hidden_size)
        
    
    def forward(self, x, month, day, hour):
        b = x.shape[0]
        s = x.shape[-2]
        x = rearrange(x, 'b t s h -> (b s) t h')
        hour = self.hour_embedding(hour)
        day = self.day_embedding(day)
        month = self.month_embedding(month)
        time = hour + day + month
        time = time.repeat_interleave(s, dim=0)
        x = x + time
        out, _ = self.net(x)
        out = rearrange(out, '(b s) t h -> b t (s h)', b=b)
        return out